package com.hound.config.datasource;

import com.alibaba.druid.pool.DruidDataSource;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * 动态数据源注册
 * Created by lupeng on 2018/4/1.
 */
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    // 默认数据源
    private DataSource defaultDataSource;

    private Map<String, Object> customDataSources = new HashMap<>();

    @Override
    public void setEnvironment(Environment environment) {
        System.out.println("DynamicDataSourceRegister.setEnvironment()");
        initDefaultDataSource(environment);
        initCustomDataSource(environment);
    }

    /**
     * 加载主数据源配置.
     *
     * @param env
     */
    private void initDefaultDataSource(Environment env) {
        // 读取主数据源
        String driver = env.getProperty("spring.datasource.driver-class-name");
        String url = env.getProperty("spring.datasource.url");
        String username = env.getProperty("spring.datasource.username");
        String password = env.getProperty("spring.datasource.password");
        Map<String, Object> dsParam = new HashMap<>();
        dsParam.put("driver", driver);
        dsParam.put("url", url);
        dsParam.put("username", username);
        dsParam.put("password", password);

        //创建数据源
        defaultDataSource = buildDataSource(dsParam);
        DataSourceContextHandler.setDataSourceId("default");
    }

    /**
     * 加载其他自定义数据源
     *
     * @param env 环境
     */
    private void initCustomDataSource(Environment env) {
//        custom:
//        datasource:
//        names: ds1
        String names = env.getProperty("custom.datasource.names").toString();
        String[] nameArr = names.split(",");
        for (String name : nameArr) {
            String driver = getPropertyFromEnv("custom.datasource." + name + ".driverClassName", env);
            String url = getPropertyFromEnv("custom.datasource." + name + ".url", env);
            String username = getPropertyFromEnv("custom.datasource." + name + ".username", env);
            String password = getPropertyFromEnv("custom.datasource." + name + ".password", env);
            // 自定义数据源
            customDataSources.put(name, buildDataSource(driver, url, username, password));
            DataSourceContextHandler.setDataSourceId(name);
        }
    }


    /**
     * 创建数据源
     *
     * @param driver   驱动
     * @param url      URL
     * @param username 用户名
     * @param password 密码
     * @return javax.sql.DataSource
     */
    private DataSource buildDataSource(String driver, String url, String username, String password) {
        Map<String, Object> map = new HashMap<>();
        map.put("driver", driver);
        map.put("url", url);
        map.put("username", username);
        map.put("password", password);
        return this.buildDataSource(map);
    }

    /**
     * 创建数据源
     *
     * @param dsParam 参数集合
     */
    private DataSource buildDataSource(Map<String, Object> dsParam) {
        String driver = dsParam.get("driver").toString();
        String url = dsParam.get("url").toString();
        String username = dsParam.get("username").toString();
        String password = dsParam.get("password").toString();

        DruidDataSource dataSource = null;
        try {
            dataSource = new DruidDataSource();
            dataSource.setDriverClassName(driver);
            dataSource.setUrl(url);
            dataSource.setUsername(username);
            dataSource.setPassword(password);
            dataSource.setFilters("stat,wall");
        } catch (Exception e){
            e.printStackTrace();
        }
        return dataSource;
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata annotationMetadata, BeanDefinitionRegistry registry) {
        System.out.println("DynamicDataSourceRegister.registerBeanDefinitions()");
        Map<Object, Object> targetDataSources = new HashMap<Object, Object>();
        targetDataSources.put("dataSource", defaultDataSource);
        targetDataSources.putAll(customDataSources);
        // 创建DynamicDataSource
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setSynthetic(true);
        MutablePropertyValues mpv = beanDefinition.getPropertyValues();
        //添加属性：AbstractRoutingDataSource.defaultTargetDataSource
        mpv.addPropertyValue("defaultTargetDataSource", defaultDataSource);
        mpv.addPropertyValue("targetDataSources", targetDataSources);
        registry.registerBeanDefinition("dataSource", beanDefinition);
    }


    private static String getPropertyFromEnv(String key, Environment env) {
        if (key != null && key.trim().length() > 0) {
            return env.getProperty(key).toString();
        }
        return null;
    }
}
